__doc__ =   """ Functions for creation of PDB files with data in the
                B-factor and Q (occupancy) columns.
            """

from Einput import einput
from PDBWriter import PDBWriter

pdb_writer = PDBWriter()

# The B and Q columns can contain only numeric values, thus any data which we wish to store
# in those columns needs to be converted. The following functions convert data to numerica form.
# boolean type can also be treated as number
def number(data, rest_val):
    try:
        return float(data)
    except TypeError:
        return rest_val

# object has to support lenght iterators and generators
# are a problem and should not be stored in the xtra dict.
def iterable(data, rest_val):
    try:
        return float(len(data))
    except TypeError:
        return rest_val


def BeQu(file_name, ents, level, b_key, b_mode =None, b_val =0.0, q_key =None, q_mode =None, q_val =0.0):
    """ Writes data into b and q columns of a PDB file.
    """

    b_mode = (b_mode or 'number')   # B
    q_mode = (q_mode or 'number')   # Q

    entities = einput(ents, level)

    for entity in entities:
        q_data = eval(q_mode)(entity.xtra.get(q_key), b_val) # occupancy
        b_data = eval(b_mode)(entity.xtra.get(b_key), q_val) # b-factor

        if level != 'A':
            atoms = einput(entity, 'A')
            for atom in atoms:
                if b_key:
                    atom.set_bfactor(b_data)
                if q_key:
                    atom.set_occupancy(q_data)
        else:
            if b_key:
                entity.set_bfactor(b_data)
            if q_key:
                entity.set_occupancy(q_data)
    pdb_writer.set_structure(ents)  # we try to preserve the headers
    pdb_writer.save(file_name)

def BeQuick(file_name, ents):
    """ Wrapeer around pdb_writer
    """
    pdb_writer.set_structure(ents)  # we try to preserve the headers
    pdb_writer.save(file_name)
    

#EOF